#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
from pathlib import Path
from .base_request import BaseHTTPClient
from typing import List, Dict
from dotenv import load_dotenv

load_dotenv(Path(__file__).parent.parent / ".env")

RERANK_URL = os.getenv("RERANK_URL")
RERANK_API_KEY = os.getenv("RERANK_API_KEY")
RERANK_MODEL = os.getenv("RERANK_MODEL")


class RerankerRequest(BaseHTTPClient):
    def __init__(self):
        super().__init__(
            base_url=RERANK_URL
        )
        self.api_key = RERANK_API_KEY

    async def rerank(
            self,
            query: str,
            documents: List[str],
            top_n: int = 5
    ) -> List[Dict]:
        return await self._request(
            "POST",
            json={
                "model": RERANK_MODEL,
                "query": query,
                "documents": documents,
                "top_n": top_n
            },
            headers={
                "Authorization": f"Bearer {RERANK_API_KEY}",
                "Content-Type": "application/json"
            }
        )
